package com.xs.langchain4j_springboot.config;

import com.xs.langchain4j_springboot.service.ToolService;
import dev.langchain4j.community.model.dashscope.QwenEmbeddingModel;
import dev.langchain4j.memory.ChatMemory;
import dev.langchain4j.memory.chat.MessageWindowChatMemory;
import dev.langchain4j.model.chat.ChatLanguageModel;
import dev.langchain4j.model.chat.StreamingChatLanguageModel;
import dev.langchain4j.rag.content.retriever.ContentRetriever;
import dev.langchain4j.rag.content.retriever.EmbeddingStoreContentRetriever;
import dev.langchain4j.service.*;
import dev.langchain4j.store.embedding.EmbeddingStore;
import dev.langchain4j.store.embedding.inmemory.InMemoryEmbeddingStore;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import reactor.core.publisher.Flux;

import java.util.stream.Stream;

@Configuration
public class AiConfig {
    public interface Assistant {
        String chat(String message);
        TokenStream stream(String message);

        @SystemMessage("""
        你是一个商品分类审核师，审核人员会告知你需要审核的商品名称，请根据商品名称判断商品是否属于哪个类目。
        注意，判断结果必须存在于从向量数据库检索到的结果中，
        如果无法检索到结果，直接回答无法判断。回答的形式为：{{message}} 属于 某某类目。
        """)
        TokenStream categoryAudit(@UserMessage String message);


        TokenStream PDFReader();


        @SystemMessage("你是一个嘲讽大师")
        @UserMessage("根据以下输入生成一段嘲讽的语句: {{text}}")
        String roastMe(@V("text") String message);

        @SystemMessage("You are a professional translator.")
        @UserMessage("Translate the following text to English: {{text}}")
        String doSomething(@V("text") String message, String en);

        @SystemMessage("""
                你是志伟航空公司的客服，请以专业的方式来回答客户的问题
                你现在正在通过在线聊天系统与客户互动
                在提供订单信息或者针对订单做操作之前，你必须从用户处获取以下信息：订单号，客户姓名。
                请说中文
                今天的日期是 {{current_date}}。
                """)
        TokenStream ticketAssistant(@UserMessage String message, @V("current_date") String currentDate);
    }

    @Bean
    public EmbeddingStore embeddingStore() {
        return new InMemoryEmbeddingStore();
    }

    @Bean
    public Assistant assistant(ChatLanguageModel qwenChatModel, StreamingChatLanguageModel qwenStreamingChatModel, ToolService toolService, QwenEmbeddingModel qwenEmbeddingModel, EmbeddingStore embeddingStore) {
        ChatMemory chatMemory = MessageWindowChatMemory.withMaxMessages(10);
        ContentRetriever contentRetriever = EmbeddingStoreContentRetriever
                .builder()
                .embeddingStore(embeddingStore)
                .embeddingModel(qwenEmbeddingModel)
                .maxResults(6)
                .minScore(0.6)
                .build();
        Assistant assistant = AiServices.builder(Assistant.class)
                .tools(toolService)
                .chatLanguageModel(qwenChatModel)
                .streamingChatLanguageModel(qwenStreamingChatModel)
                .chatMemory(chatMemory)
                .contentRetriever(contentRetriever)
                .build();
        return assistant;
    }

    public interface AssistantUnique {
        String chat(@MemoryId int memoryId, @UserMessage String userMessage);

        String chat(String message);

        TokenStream stream(@MemoryId int memoryId, @UserMessage String userMessage);
    }

    @Bean
    public AssistantUnique assistantUnique(ChatLanguageModel qwenChatModel, StreamingChatLanguageModel qwenStreamingChatModel) {

        AssistantUnique assistantUnique = AiServices.builder(AssistantUnique.class)
                .chatLanguageModel(qwenChatModel)
                .streamingChatLanguageModel(qwenStreamingChatModel)
                .chatMemoryProvider(memoryId ->
                        MessageWindowChatMemory.builder().maxMessages(10)
                                .id(memoryId).build()
                )
                .build();
        return assistantUnique;
    }
}
